package org.demo.test.springbatch;

import com.hundsun.jrescloud.batch.mq.constant.Queues;
import com.hundsun.jrescloud.batch.mq.util.RabbitUtil;
import org.springframework.amqp.core.Queue;
import org.springframework.amqp.rabbit.AsyncRabbitTemplate;
import org.springframework.amqp.rabbit.connection.ConnectionFactory;
import org.springframework.batch.integration.partition.MessageChannelPartitionHandler;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.integration.amqp.outbound.AsyncAmqpOutboundGateway;
import org.springframework.integration.annotation.ServiceActivator;
import org.springframework.integration.channel.DirectChannel;
import org.springframework.integration.channel.QueueChannel;
import org.springframework.integration.config.AggregatorFactoryBean;
import org.springframework.integration.core.MessagingTemplate;
import org.springframework.integration.scheduling.PollerMetadata;
import org.springframework.scheduling.support.PeriodicTrigger;
import org.springframework.util.ErrorHandler;

@Configuration
public class ClusterJobMasterAutoConfig {

    @Autowired
    @Qualifier("batchConnectionFactory")
    private ConnectionFactory connectionFactory;

    @Autowired
    RabbitUtil rabbitUtil;

    /**
     * 集群任务分片队列
     *
     * @return
     */
    @Bean
    public Queue jobSplit() {
        return new Queue(Queues.JOB_SPLIT, true, false, false);
    }


    @Bean("masterPartitionHandler")
    public MessageChannelPartitionHandler partitionHandler() {
        MessageChannelPartitionHandler handler = new MessageChannelPartitionHandler();
        handler.setStepName("");
//        handler.setGridSize(5);
        MessagingTemplate template = new MessagingTemplate();
        // 设置发送任务的Channel
        template.setDefaultChannel(masterRequestChannel());
        handler.setMessagingOperations(template);
        // 接受Slave返回信息的Channel，注意这是一个AggregatedChannel。因为Master需要把所有Slave返回的任务信息整合在一起，才能判断任务是否成功完成。
        handler.setReplyChannel(aggregatedChannel());
        return handler;
    }

//    @Bean
//    @ServiceActivator(inputChannel = "masterRequestChannel")
//    public AmqpOutboundEndpoint requestEndpoint(@Qualifier("batchRabbitTemplate") RabbitTemplate template) {
//        template.setReplyTimeout(-1);
//        template.setReplyAddress("mysplitQueue");
//        AmqpOutboundEndpoint endpoint = new AmqpOutboundEndpoint(template);
//        // 调用 rabbitTemplate sendAndReceive()
//        endpoint.setExpectReply(false);
//        endpoint.setOutputChannel(masterReplyChannel());
//        endpoint.setRoutingKey(Queues.JOB_SPLIT);
//        return endpoint;
//    }

    @Bean
    @ServiceActivator(inputChannel = "masterRequestChannel")
    public AsyncAmqpOutboundGateway amqpOutbound() {
        //添加匿名回调队列
        Queue replyQueue = rabbitUtil.addQueue();
        AsyncRabbitTemplate asyncRabbitTemplate = new AsyncRabbitTemplate(connectionFactory, null, Queues.JOB_SPLIT, replyQueue.getName());
        asyncRabbitTemplate.setReceiveTimeout(-1);
        asyncRabbitTemplate.setEnableConfirms(true);
        AsyncAmqpOutboundGateway outbound = new AsyncAmqpOutboundGateway(asyncRabbitTemplate);
        outbound.setOutputChannel(masterReplyChannel());
        return outbound;
    }


    @Bean
    @ServiceActivator(inputChannel = "masterReplyChannel")
    public AggregatorFactoryBean aggregator() {
        AggregatorFactoryBean aggregatorFactoryBean = new AggregatorFactoryBean();
        aggregatorFactoryBean.setProcessorBean(partitionHandler());
//        aggregatorFactoryBean.setCountsEnabled(false);
        aggregatorFactoryBean.setOutputChannel(aggregatedChannel());
        return aggregatorFactoryBean;
    }

    @Bean("masterRequestChannel")
    public DirectChannel masterRequestChannel() {
        return new DirectChannel();
    }

    @Bean("masterReplyChannel")
    public QueueChannel masterReplyChannel() {
        return new QueueChannel();
    }

    @Bean
    public QueueChannel aggregatedChannel() {
        return new QueueChannel();
    }

    @Bean(name = PollerMetadata.DEFAULT_POLLER)
    public PollerMetadata defaultPoller() {
        PollerMetadata pollerMetadata = new PollerMetadata();
        pollerMetadata.setTrigger(new PeriodicTrigger(1 * 1000));
        pollerMetadata.setErrorHandler(new ErrorHandler() {
            @Override
            public void handleError(Throwable cause) {
                throw new RuntimeException(cause);
            }
        });
        return pollerMetadata;
    }
}
